
//#include <cstdlib>
#include <unistd.h>
#include <fcntl.h>
#include <sys/mman.h> // Include this for mprotect
//#include <string>     // Include for std::string
#include <android/log.h>
#include <string.h>   // Include for memcpy and other C-style string functions

#include "zygisk.hpp"

using zygisk::Api;
using zygisk::AppSpecializeArgs;
using zygisk::ServerSpecializeArgs;

#define LOGD(...) __android_log_print(ANDROID_LOG_DEBUG, "MyModule", __VA_ARGS__)
#define LOGE(...) __android_log_print(ANDROID_LOG_ERROR, "MyModule", __VA_ARGS__)

class MyModule : public zygisk::ModuleBase
{
public:
    void onLoad(Api *api, JNIEnv *env) override
    {
        LOGD("onload from my zygisk module!!!!");
        this->api = api;
        this->env = env;
    }

    // void preAppSpecialize(AppSpecializeArgs *args) override {
    //     // Use JNI to fetch our process name
    //     const char *process = env->GetStringUTFChars(args->nice_name, nullptr);
    //     preSpecialize(process);
    //     env->ReleaseStringUTFChars(args->nice_name, process);
    // }

    // void preServerSpecialize(ServerSpecializeArgs *args) override {
    //     preSpecialize("system_server");
    // }
    void postAppSpecialize(const AppSpecializeArgs *args) override
    {
        this->loadDex(env, "assets/classes.dex");
        this->replaceArtMethod(env);
    }

private:
    Api *api;
    JNIEnv *env;

    // void preSpecialize(const char *process) {
    //     // Demonstrate connecting to to companion process
    //     // We ask the companion for a random number
    //     unsigned r = 0;
    //     int fd = api->connectCompanion();
    //     read(fd, &r, sizeof(r));
    //     close(fd);
    //     LOGD("process=[%s], r=[%u]\n", process, r);

    //     // Since we do not hook any functions, we should let Zygisk dlclose ourselves
    //     api->setOption(zygisk::Option::DLCLOSE_MODULE_LIBRARY);
    // }

    void loadDex(JNIEnv *env, const char *dex_path)
    {
        // 1. Get the current application's ClassLoader
        jclass application_class = env->FindClass("android/app/Application");
        jmethodID get_class_loader_method = env->GetMethodID(application_class, "getClassLoader", "()Ljava/lang/ClassLoader;");
        jobject app_context = env->CallObjectMethod(application_class, get_class_loader_method);

        // 2. Find the DexClassLoader class
        jclass dex_class_loader_class = env->FindClass("dalvik/system/DexClassLoader");
        jmethodID dex_class_loader_ctor = env->GetMethodID(dex_class_loader_class, "<init>", "(Ljava/lang/String;Ljava/lang/String;Ljava/lang/String;Ljava/lang/ClassLoader;)V");

        // 3. Create a DexClassLoader instance
        jstring dex_path_jstr = env->NewStringUTF(dex_path);
        jstring optimized_dir_jstr = env->NewStringUTF(nullptr);
        jobject dex_class_loader = env->NewObject(dex_class_loader_class, dex_class_loader_ctor, dex_path_jstr, optimized_dir_jstr, nullptr, app_context);

        // 4. Find the loadClass method
        jmethodID load_class_method = env->GetMethodID(dex_class_loader_class, "loadClass", "(Ljava/lang/String;)Ljava/lang/Class;");

        // 5. Load your class from the DEX file
        jstring fake_location_class_name = env->NewStringUTF("org.eu.laremehpe.FakeLocation");
        jclass fake_location_class = (jclass)env->CallObjectMethod(dex_class_loader, load_class_method, fake_location_class_name);

        if (fake_location_class == nullptr)
        {
            // Handle error: class not found
            LOGE("Failed to load FakeLocation class from DEX file.");
        }
        else
        {
            LOGD("Successfully loaded FakeLocation class.");
        }

        // You can now use fake_location_class to find and hook methods
    }
    size_t getArtMethodSize(JNIEnv *env)
    {
        jclass klz = env->FindClass("org/eu/laremehpe/FakeLocation");
        jmethodID m1 = env->GetStaticMethodID(klz, "isFromMockProvider", "()Z");
        jmethodID m2 = env->GetStaticMethodID(klz, "<clinit>", "()V"); // 任何相邻方法
        return reinterpret_cast<size_t>(m2) - reinterpret_cast<size_t>(m1);
    }
    jobject getReflectedMethod(JNIEnv *env, const char *cls,
                               const char *name, const char *sig)
    {
        jclass clazz = env->FindClass(cls);
        jmethodID mid = env->GetMethodID(clazz, name, sig);
        return env->ToReflectedMethod(clazz, mid, JNI_FALSE);
    }
    void *getArtMethod(JNIEnv *env, jobject method)
    {
        return reinterpret_cast<void *>(env->FromReflectedMethod(method));
    }
    void replaceArtMethod(JNIEnv *env)
    {
        jobject src = getReflectedMethod(env,
                                         "android/location/Location", "isFromMockProvider", "()Z");
        jobject dst = getReflectedMethod(env,
                                         "org/eu/laremehpe/FakeLocation", "isFromMockProvider", "()Z");

        void *smeth = getArtMethod(env, src);
        void *dmeth = getArtMethod(env, dst);
        size_t size = getArtMethodSize(env);

        // 解锁页保护 & 拷贝
        uintptr_t page = reinterpret_cast<uintptr_t>(smeth) & ~0xfff;
        mprotect((void *)page, 0x2000, PROT_READ | PROT_WRITE | PROT_EXEC);
        memcpy(smeth, dmeth, size);
        mprotect((void *)page, 0x2000, PROT_READ | PROT_EXEC);
    }
};

// static int urandom = -1;

// static void companion_handler(int i) {
//     if (urandom < 0) {
//         urandom = open("/dev/urandom", O_RDONLY);
//     }
//     unsigned r;
//     read(urandom, &r, sizeof(r));
//     LOGD("companion r=[%u]\n", r);
//     write(i, &r, sizeof(r));
// }

// Register our module class and the companion handler function
REGISTER_ZYGISK_MODULE(MyModule)
// REGISTER_ZYGISK_COMPANION(companion_handler)
